#!/usr/bin/env python
"""
This is a small example script on how to use the shtools python wrapper to plot a map from
random spherical harmonic coefficients. It requires the modules numpy, matplotlib and basemap 
for the plotting procedures

by Matthias Meschede (2013)
"""
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, shiftgrid
import numpy as np
import sys
import shtools
import ipdb

def main():
    test_correlation()

def test1():
    lmax = 80
    coeffs = np.random.rand(2*(lmax+1)*(lmax+1)).reshape(2,lmax+1,lmax+1) - 0.5
    data = shtools.pymakegriddh(coeffs) #nlat is optional
    plot_grid(data)
    plt.show()

def plot_grid(data):
    nlat = data.shape[0]
    interval = 180.0/nlat
    fig = plt.figure()
    m = Basemap(llcrnrlon=-180.,llcrnrlat=-90,urcrnrlon=180.,urcrnrlat=90.,\
                    resolution='c',area_thresh=10000.,projection='cyl')

    lats = np.arange(-90.+interval/2.,90.0,interval)
    lons = np.arange(0.+interval/2.,360.,interval)
    nx,ny = 720,360
    topoin,lons = shiftgrid(180.,-data[:,:],lons,start=False)
    scalarmap = m.transform_scalar(topoin,lons,lats,nx,ny)
    im = m.imshow(scalarmap)
    m.colorbar(im)
    m.drawcoastlines()

def test_correlation():
    depth = 250
    sys.path.append('/home/matthias/projects/python/tomographic_models/scripts/regional_models')
    from regional_models import get_regional_model #for regional models
    print 'analyzing the fichtneu model'
    model_fichtneu = get_regional_model('fichtneu')
    lmaxt_fichtneu = 20        #this is the bandwith of the local filter
    k_fichtneu = 3            #this is the number of tapers used
    nlat_fichtneu, clat_fichtneu, clon_fichtneu, theta0_fichtneu_deg = model_fichtneu.get_taper_params()
    lmax_fichtneu = nlat_fichtneu/2 - 1
    ls_fichtneu = np.arange(0,lmax_fichtneu - lmaxt_fichtneu + 1)
    grid = model_fichtneu.get_global_grid(depth,nlat=nlat_fichtneu)
    print 'spherical harmonics expansion...'
    coeffs_fichtneu_rot = shtools.pyshexpanddh(grid)
    coeffs_fichtneu = shtools.pyshrotaterealcoef(coeffs_fichtneu_rot,np.radians([0.,57.5,0.]))
    print 'expanded grid to coeff shape:',coeffs_fichtneu.shape
    print 'calculating multitaper regional spectral estimate...'
    #compute spherical cap taper values
    theta0_fichtneu = np.radians(theta0_fichtneu_deg)
    #theta0_fichtneu = theta0_dna10
    print 'getting tapers for direct bias calculation...'
    tapers_fichtneu, eigenvalues_fichtneu, taper_order_fichtneu = shtools.pyshreturntapers(theta0_fichtneu, lmaxt_fichtneu)
    print 'the first %d tapers have concentration factors:'%k_fichtneu
    print eigenvalues_fichtneu[:k_fichtneu]
    print 'using cap with radius of %.0f degrees'%theta0_fichtneu_deg
    psd1d_fichtneu,dev_fichtneu = shtools.pyshmultitaperse(coeffs_fichtneu, tapers_fichtneu, taper_order_fichtneu,
                                            k_fichtneu, theta0_fichtneu, clat_fichtneu, clon_fichtneu)

    fichtneu_vs_fichtneu = \
            shtools.pyshlocalizedadmitcorr(tapers_fichtneu[:,:k_fichtneu],
                                      taper_order_fichtneu[:k_fichtneu],
                                      clat_fichtneu, clon_fichtneu,
                                      coeffs_fichtneu, coeffs_fichtneu)
    ipdb.set_trace()
    #---- plotting ----
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.set_xscale('log')
    ax.set_yscale('log')
    lmin_fichtneu = max(lmaxt_fichtneu, int(2.*360./(2.*theta0_fichtneu_deg)))
    ax.plot(ls_fichtneu[lmin_fichtneu:], psd1d_fichtneu[lmin_fichtneu:],label=r'Fichtneu, %d$^\circ$ cap radius'%theta0_fichtneu_deg)
    plt.show()
    
if __name__ == "__main__":
    main()
